import os
import easyocr
import pandas as pd
from PIL import Image
import re
from datetime import datetime

def extract_keywords(text):
    pattern = r'(?:1[3-9]\d{9}|[a-zA-Z0-9_-]{6,20})'
    keywords = re.findall(pattern, text)
    return list(dict.fromkeys(keywords))  # 去重

def easyocr_nested_folders(root_folder, formatted_date, batch_size=100):
    """
    使用EasyOCR递归识别两层文件夹中的图片文字
    :param root_folder: 根文件夹路径
    :param batch_size: 每次保存的批量大小
    """
    supported_formats = ('.png', '.jpg', '.jpeg', '.bmp', '.tiff', '.gif')
    reader = easyocr.Reader(['ch_sim', 'en'])
    results = []
    total_processed = 0  # 计数已处理的文件

    # 遍历根文件夹
    for dir1 in os.listdir(root_folder):
        dir1_path = os.path.join(root_folder, dir1)

        if not os.path.isdir(dir1_path):
            continue

        # 遍历第二层文件夹
        for filename in os.listdir(dir1_path):
            filepath = os.path.join(dir1_path, filename)

            try:
                # 使用OCR识别文字
                img = Image.open(filepath)
                ocr_results = reader.readtext(img, detail=0)
                info = ' '.join(ocr_results[1:len(ocr_results) - 1])

                keywords = extract_keywords(info)
                results.append({
                    'keyword': dir1,
                    "name": ocr_results[0],
                    "info": info,
                    'wechat': '  '.join(keywords)
                })

                total_processed += 1

                # 如果达到批量大小，保存数据并清空结果列表
                if total_processed % batch_size == 0:
                    df = pd.DataFrame(results)
                    filename_batch = f"微信视频号数据采集_{formatted_date}_batch_{total_processed // batch_size}.xlsx"
                    df.to_excel(filename_batch, index=False)
                    print(f"OCR results saved to {filename_batch}")
                    results.clear()  # 清空结果列表，释放内存

            except Exception as e:
                print(f"Error processing {dir1}/{filename}: {str(e)}")
                results.append({
                    'folder_level1': dir1,
                    'text': f'ERROR: {str(e)}'
                })

    # 处理剩余的结果
    if results:
        df = pd.DataFrame(results)
        filename = f"微信视频号数据采集_{formatted_date}_final.xlsx"
        df.to_excel(filename, index=False)
        print(f"OCR results saved to {filename}")

now = datetime.now()
formatted_date = now.strftime("%Y_%m_%d")
relative_path = os.path.join(formatted_date)
easyocr_nested_folders(relative_path, formatted_date)
